# validate_models.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
VEHICLE_MODEL_PATH = r"D:\CodeCNN\intel\vehicle-detection-0202\FP16\vehicle-detection-0202.xml"
ROAD_MODEL_PATH = r"D:\CodeCNN\intel\road-segmentation-adas-0001\FP16\road-segmentation-adas-0001.xml"
IMAGE_PATH = r"H:\xiaomi\test\192.168.1.64_01_20251027093805446_frame_000061.png"  # 建议使用含车辆和道路的图像

# -----------------------------
# 初始化 OpenVINO
# -----------------------------
core = Core()
device = "CPU"  # 可选 "GPU"（若支持）

# -----------------------------
# 1. 加载车辆检测模型
# -----------------------------
print("Loading vehicle detection model...")
vehicle_model = core.read_model(model=VEHICLE_MODEL_PATH)
vehicle_compiled = core.compile_model(model=vehicle_model, device_name=device)
vehicle_input_layer = vehicle_compiled.input(0)
vehicle_output_layer = vehicle_compiled.output(0)

# 获取输入尺寸
_, _, h_det, w_det = vehicle_input_layer.shape
print(f"Vehicle detection input shape: {h_det}x{w_det}")

# -----------------------------
# 2. 加载道路分割模型
# -----------------------------
print("Loading road segmentation model...")
road_model = core.read_model(model=ROAD_MODEL_PATH)
road_compiled = core.compile_model(model=road_model, device_name=device)
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

_, _, h_seg, w_seg = road_input_layer.shape
print(f"Road segmentation input shape: {h_seg}x{w_seg}")

# -----------------------------
# 3. 读取并预处理图像
# -----------------------------
image = cv2.imread(IMAGE_PATH)
if image is None:
    raise FileNotFoundError(f"Cannot load image: {IMAGE_PATH}")

orig_h, orig_w = image.shape[:2]
print(f"Original image size: {orig_w}x{orig_h}")

# -----------------------------
# 4. 车辆检测推理
# -----------------------------
# 调整图像尺寸并归一化（模型通常期望 BGR + [0,1] 或 [0,255]，此模型用 [0,255]）
resized_det = cv2.resize(image, (w_det, h_det))
input_det = np.expand_dims(resized_det.transpose(2, 0, 1), axis=0)  # HWC -> NCHW

# 推理
detections = vehicle_compiled([input_det])[vehicle_output_layer]
# detections shape: [1, 1, N, 7] → [image_id, label, conf, x_min, y_min, x_max, y_max]

# -----------------------------
# 5. 道路分割推理
# -----------------------------
resized_seg = cv2.resize(image, (w_seg, h_seg))
input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)

segmentation = road_compiled([input_seg])[road_output_layer]
# segmentation shape: [1, 1, H, W] → 每个像素是 0（非道路）或 1（道路）

seg_mask = segmentation[0, 0]  # HxW
seg_mask = (seg_mask < 0.5).astype(np.uint8) * 255  # 二值化并转为 0-255

# 将分割掩码缩放回原图尺寸
seg_mask_resized = cv2.resize(seg_mask, (orig_w, orig_h))

# 创建彩色掩码（绿色道路）
color_mask = np.zeros_like(image)
color_mask[:, :, 1] = seg_mask_resized  # Green channel

# -----------------------------
# 6. 可视化结果
# -----------------------------
output_image = image.copy()

# 绘制车辆检测框（置信度 > 0.5）
for det in detections[0, 0]:
    conf = float(det[2])
    if conf > 0.5:
        x_min = int(det[3] * orig_w)
        y_min = int(det[4] * orig_h)
        x_max = int(det[5] * orig_w)
        y_max = int(det[6] * orig_h)
        cv2.rectangle(output_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.putText(output_image, f"Car: {conf:.2f}", (x_min, y_min - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

# 叠加道路分割（半透明绿色）
overlay = cv2.addWeighted(output_image, 0.7, color_mask, 0.3, 0)


# -----------------------------
# 7. 等比例缩放并显示（目标窗口：1360x768）
# -----------------------------
def resize_to_fit(image, target_width=1360, target_height=768):
    """
    等比例缩放图像到目标尺寸内，不足部分用黑色填充（letterbox）
    """
    h, w = image.shape[:2]
    scale = min(target_width / w, target_height / h)
    new_w = int(w * scale)
    new_h = int(h * scale)

    resized = cv2.resize(image, (new_w, new_h))

    # 创建目标尺寸的黑色画布
    canvas = np.zeros((target_height, target_width, 3), dtype=np.uint8)

    # 居中放置缩放后的图像
    y_offset = (target_height - new_h) // 2
    x_offset = (target_width - new_w) // 2
    canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized

    return canvas


# 生成带红路+检测框的结果图（overlay）
# ...（前面的推理和overlay生成代码保持不变）...

# 缩放到1360x768并保持比例
display_image = resize_to_fit(overlay, 1360, 768)

# 创建可调整大小的窗口（但固定为1360x768显示）
cv2.namedWindow("Vehicle Detection + Road Segmentation", cv2.WINDOW_NORMAL)
cv2.resizeWindow("Vehicle Detection + Road Segmentation", 1360, 768)
cv2.imshow("Vehicle Detection + Road Segmentation", display_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 从原图路径生成结果名，例如 input.jpg → input_result.jpg
input_name = os.path.splitext(os.path.basename(IMAGE_PATH))[0]
output_path = f"{input_name}_result.jpg"

# 如果已存在，则加时间戳避免覆盖
if os.path.exists(output_path):
    timestamp = time.strftime("%H%M%S")
    output_path = f"{input_name}_result_{timestamp}.jpg"

cv2.imwrite(output_path, overlay)
print(f"✅ 结果已保存为: {output_path}")

print("✅ 模型验证完成！")